import os

import numpy as np
import matplotlib.pyplot as plt
from scipy.constants import physical_constants

GRAVITY = physical_constants["standard acceleration of gravity"][0]
WORKSPACE = os.getcwd()


class FinalCar:
    car_M = 0.
    car_M_c1 = 0.
    car_M_c2 = 0.
    car_theta_1 = 0.
    car_theta_2 = 0.
    car_l_c1 = 0.
    car_l_c2 = 0.
    car_I_c1 = 0.
    car_I_c2 = 0.
    car_name = ""
    result_data_time = list()
    result_data_position = list()
    result_data_theta_1 = list()
    result_data_theta_2 = list()

    def __init__(self, properties: tuple) -> None:
        self.setProperties(properties)
        
    def setProperties(self, properties: tuple) -> None:
        self.car_M, M_1, M_2, m_1, m_2, theta_1, theta_2, l_1, l_2 = properties
        self.car_theta_1 = self.degreeConvert(theta_1)
        self.car_theta_2 = self.degreeConvert(theta_2)
        self.car_M_c1 = M_1 + m_1
        self.car_M_c2 = M_2 + m_2
        self.car_l_c1 = (M_1 + m_1 / 2) * l_1 / self.car_M_c1
        self.car_l_c2 = (M_2 + m_2 / 2) * l_2 / self.car_M_c2
        
        def inertia(m: float, M: float, l:float):
            return m * np.square(l / (m + M)) * (M * m / 4 + np.square(M + m / 2) / 3)
        
        self.car_I_c1 = inertia(m_1, M_1, l_1)
        self.car_I_c2 = inertia(m_2, M_2, l_2)
        self.car_name = "M-%s_M1-%s_M2-%s_m1-%s_m2-%s_t1-%s_t2-%s_l1-%s_l2-%s"%properties
        
    def calculatePositionAcceleration(self, theta_1: float, theta_2: float) -> float:
        above = 0
        below = self.car_M
        for index, theta in enumerate([theta_1, theta_2]):
            l_c = getattr(self, "car_l_c{:d}".format(index + 1))
            I_c = getattr(self, "car_I_c{:d}".format(index + 1))
            M_c = getattr(self, "car_M_c{:d}".format(index + 1))
            public = np.square(l_c * M_c) / (I_c + np.square(l_c) * M_c)
            above += public * GRAVITY * theta
            below += M_c - public
        return above / below
            
    def calculateAngleAcceleration(self, acceleration: float, theta: float, index: str) -> float:
        l_c = getattr(self, "car_l_c" + index)
        I_c = getattr(self, "car_I_c" + index)
        M_c = getattr(self, "car_M_c" + index)
        return -l_c * M_c * (acceleration + GRAVITY * theta) / (I_c + np.square(l_c) * M_c)
    
    def simulate(self, numbers=int(1e4), dt=1e-3, saveData=True) -> None:
        self.result_data_time = list()
        self.result_data_position = list()
        self.result_data_theta_1 = list()
        self.result_data_theta_2 = list()
        position = 0.
        velocity = 0.
        theta_1 = self.car_theta_1
        theta_1_velocity = 0.
        theta_2 = self.car_theta_2
        theta_2_velocity = 0.
        for time in np.linspace(0, numbers * dt, numbers + 1):
            self.result_data_time.append(time)
            self.result_data_position.append(position)
            self.result_data_theta_1.append(theta_1)
            self.result_data_theta_2.append(theta_2)
            
            acceleration = self.calculatePositionAcceleration(theta_1, theta_2)
            velocity += acceleration * dt
            position += velocity * dt
            
            theta_1_acceleration = self.calculateAngleAcceleration(acceleration, theta_1, '1')
            theta_1_velocity += theta_1_acceleration * dt
            theta_1 += theta_1_velocity * dt
            
            theta_2_acceleration = self.calculateAngleAcceleration(acceleration, theta_2, '2')
            theta_2_velocity += theta_2_acceleration * dt
            theta_2 += theta_2_velocity * dt
        np.savetxt(os.path.join(WORKSPACE, self.car_name + ".txt"), tuple(zip(self.result_data_time, self.result_data_theta_1, self.result_data_theta_2, self.result_data_position)))
        self.plotPicture(numbers, dt)
            
    def plotPicture(self, numbers: int, dt: float):
        plt.clf()
        plt.figure(figsize=(4, 9))
        
        plt.subplot(3, 1, 1)
        plt.title("$\\theta_1\sim t$")
        plt.plot(self.result_data_time, self.result_data_theta_1)
        plt.ylabel("$\\theta_1$")
        plt.grid()
        
        plt.subplot(3, 1, 2)
        plt.title("$\\theta_2\sim t$")
        plt.plot(self.result_data_time, self.result_data_theta_2)
        plt.ylabel("$\\theta_2$")
        plt.grid()
        
        plt.subplot(3, 1, 3)
        plt.title("$x\sim t$")
        plt.plot(self.result_data_time, self.result_data_position)
        plt.ylabel("$x/m$")
        plt.grid()
        
        plt.xlabel("$t/s$")
        plt.tight_layout()
        plt.savefig(os.path.join(WORKSPACE, self.car_name + "_N-%s_dt-%s.png"%(numbers, dt)))
        
    @staticmethod
    def degreeConvert(angle: float) -> float:
        return angle / 180 * np.pi
        
if __name__ == "__main__":
    properties = [
        (20, 10, 10, 0, 0, 15, 15, 3, 3),
        (20, 11, 1, 9, 0, 15, 15, 3, 3),
        (20, 11, 1, 9, 0, 15, 15, 5, 0.5),
    ]
    fc = FinalCar(properties[0])
    fc.simulate()
    fc.setProperties(properties[1])
    fc.simulate()
    fc.setProperties(properties[2])
    fc.simulate()